#!/usr/bin/env python
# -*- coding: UTF-8 -*-

import sys, os, time
import logging
from subprocess import Popen, PIPE
import ConfigParser
from psycopg2.pool import SimpleConnectionPool
import base64
import pickle
import traceback

class SNMPTrapSender:
    def __init__(self):
        self.trap_enable = 1  # 0 or 1
        self.trap_address = "10.64.70.122"
        self.trap_version = 2  # 2 or 3
        #for v1/v2c
        self.trap_community = "public"
        #for v3
        self.security = ""
        self.auth_name = ""
        self.auth_pass = ""
        self.priv_pass = ""

    def send_trap_message(self, type_oid, msg_oid, message):
        #print type_oid, msg_oid, message
        if(self.trap_enable == 1):
            if(self.trap_version ==2):
                self._send_snmp_trap_v2(type_oid, msg_oid, message)
            elif(self.trap_version ==3):
                self._send_snmp_trap_v3(type_oid, msg_oid, message)

    def _send_snmp_trap_v2(self, type_oid, msg_oid, message):
        cmd = "snmptrap -v 2c -c "+ self.trap_community.strip() + " " + self.trap_address.strip() + " \"\" "
        cmd = cmd + type_oid + " " + msg_oid + " s \"" + message.strip() + "\""
        os.system(cmd)

    def _send_snmp_trap_v3(self, type_oid, msg_oid, message):
        cmd = "snmptrap -v 3 -l "+ self.security +" -u "+ self.auth_name.strip() + " "
        if(self.security == "authNoPriv"):
            cmd = cmd + "-a MD5 -A " + self.auth_pass.strip() + " "
        if(self.security == "authPriv"):
            cmd = cmd + "-a MD5 -A " + self.auth_pass.strip() + " -x DES -X " + self.priv_pass.strip() + " "
        cmd = cmd + self.trap_address.strip() + " \"\" " + type_oid + " " + msg_oid + " s \"" + message.strip() + "\""
        os.system(cmd)

SNMP_PARENT_OID = ".1.3.6.1.4.1.6101.3005"
SNMP_TRAP_TYPE_PARENT_OID = SNMP_PARENT_OID + ".2.1.0"
SNMP_TRAP_MSG_PARENT_OID = SNMP_PARENT_OID + ".2.1.1"
RULE_ID_MAPPING_TO_OID ={
    1  : 6,
    2  : 7,
    3  : 4,
    4  : 5,
    5  : 1,
    6  : 8,
    7  : 9,
    8  : 10,
    11 : 2,
    12 : 3,
    13 : 11,
    14 : 12,
    16 : 14,
    17 : 15
}

def make_oid_by_rule_id(rule_id):
    type_oid = SNMP_TRAP_TYPE_PARENT_OID + "." + str(RULE_ID_MAPPING_TO_OID[rule_id])
    msg_oid = SNMP_TRAP_MSG_PARENT_OID + "." + str(RULE_ID_MAPPING_TO_OID[rule_id]) + ".0"
    return type_oid, msg_oid

def help():
    print ""
    print "Usage:"
    print "     python snmp_trap.py alert_rule_id message_content       # Notify Alert by SNMP Trap."
    print ""

logging.basicConfig(filename='/opt/TrendMicro/DTAS/ManagementServer/bin/a.log',level=logging.DEBUG)
logger = logging.getLogger("")

class TrapException(Exception):
    pass

def sendTrap(rule_id,msg,subject="DDAnMaxMTUmsg"):
    '''
    Send Trap by ruleID
    :param rule_id:
    :param msg:Trap content
    :return:
    '''
    type_oid, msg_oid = make_oid_by_rule_id(int(rule_id))
    sender = SNMPTrapSender()
    #print type_oid, msg_oid, sys.argv[2]
    #Because the content may length than MTU(1500 octets)
    logger.debug('Trap lenght:%d'%len(msg))
    if len(msg) < 1500:
        sender.send_trap_message(type_oid, msg_oid, msg)
    else:
        sender.send_trap_message(type_oid, msg_oid, subject)
if __name__=="__main__":


    if len(sys.argv) == 2:
        logger.debug(sys.argv)
        request = None
        nodeInfo = None
        msg = None
        subject = None
        rule_id = None
        try:
            request=pickle.loads(base64.b64decode(sys.argv[1]))
            logger.debug('Alert event:%s'%str(request))
            rule_id = request['rule']['rule_id']
            nodeInfo = request['fillers']['ApplianceName']+"|"
            nodeInfo = request['fillers']['ApplianceIP']
            msg = request['formats']['message']
            subject = request['formats']['subject']

        except Exception as e:
            logger.error(traceback.format_exc())
        else:
            logger.debug('Rule_id:%s'%rule_id)
            logger.debug('NodeInfo:%s'%nodeInfo)
            logger.debug('Subject:%s'%subject)
            logger.debug('EmailContent:%s'%msg)
            sendTrap(rule_id,msg,subject)
        finally:
            pass
        #logger.debug('Alert event: %s'%str(alertData))
    elif len(sys.argv) == 3:
        #Valid Alert Rule ID
        if(int(sys.argv[1]) in RULE_ID_MAPPING_TO_OID):
            sendTrap(sys.argv[1],sys.argv[2])
    else:
        help()




